from model.Ark import ArkModel

import time
from langchain.tools import tool, ToolRuntime
from langchain.agents.middleware import before_model, before_agent, dynamic_prompt, ModelRequest
from langchain.agents import AgentState, create_agent
from langgraph.runtime import Runtime
from langchain.messages import RemoveMessage
from langgraph.graph.message import REMOVE_ALL_MESSAGES

""" 定义一个 agentState，这个state 中的内容，在任意的模型调用生命周期及工具的调用中都可以获取到 """
class CustomerSchema(AgentState):
    customer_field: dict

@tool
def search_my_info(runtime: ToolRuntime):
    """ 查询用户喜欢的小动物 """
    ## 由于在下面定义了 agent state 的schema，因此这里是可以直接获取到这个schema 的
    # print(runtime) runtime 中会包含 CustomerSchema 中的自定义状态

    return '小猫'

@before_agent
def before_agent_call(state: CustomerSchema, runtime: Runtime):
    ## 由于在下面定义了 agent state 的schema，因此这里是可以直接获取到这个schema 的
    print(state['customer_field'])

    ## 两种方式更改消息，不会印象 agentState
    '''
        ## 1. 直接修改数据消息本身
        state['messages'][0].content = '我最喜欢的小动物是什么?'
        
        ## 2. 通过返回新的消息格式来修改
        return {
            'messages': [
                RemoveMessage(id=REMOVE_ALL_MESSAGES),
                { 'role': 'user', 'content': '我最喜欢的小动物是什么?' }
            ]
        }
        
        备注：不论通过何种方式，只要返回了值，则都会重新开始触发 before_model_call 至下面的中间件生命周期
    '''
    return None

@before_model
def before_model_call(state: CustomerSchema, runtime: Runtime):
    print('---> before model call', time.time())
    return None

## dynamic_prompt 基于 wrap_model_call 的生命周期创建，执行顺序与 model call 一致
@dynamic_prompt
def dynamic_prompt(request: ModelRequest):
    print('123', request.messages)
    pass

def run():
    agent = create_agent(
        model=ArkModel().model,
        tools=[search_my_info],
        state_schema=CustomerSchema,
        middleware=[before_agent_call, before_model_call, dynamic_prompt]
    )

    response = agent.invoke({
        "messages": [{ 'role': 'user', 'content': '深圳的天气今天怎么样?' }],
        'customer_field': { 'aa': 123, 'bb': '456' }
    })
    print(response['messages'][-1].content)